import click
import numpy as np
import pandas as pd
from subway_qa.adapters import embedding
from subway_qa.db import build_index
from tqdm import tqdm


@click.command()
@click.argument("excel_path", type=click.Path(exists=True))
@click.option(
    "--index-path", default="qa_index.faiss", help="Path to save the FAISS index"
)
@click.option("--batch-size", default=8, help="Batch size for processing embeddings")
def sync_data(excel_path, index_path, batch_size):
    """
    Read QA pairs from Excel file, compute embeddings, and index them.
    """
    click.echo(f"Loading QA pairs from {excel_path}")

    # Read Excel file
    df = pd.read_excel(excel_path)

    # Assuming the first column is questions and the second is answers
    questions = df.iloc[:, 0].tolist()
    answers = df.iloc[:, 1].tolist()

    click.echo(f"Found {len(questions)} QA pairs")

    # Prepare QA pairs
    qa_pairs = [{"question": q, "answer": a} for q, a in zip(questions, answers)]

    # Compute embeddings in batches
    click.echo("Computing embeddings...")
    all_embeddings = []

    for question in tqdm(questions):
        all_embeddings.append(embedding(question))

    # Concatenate all embeddings
    embeddings = np.array(all_embeddings)

    click.echo(f"Generated embeddings with shape {embeddings.shape}")

    # Index the QA pairs and embeddings
    build_index(qa_pairs, embeddings, index_path)

    click.echo(f"Successfully indexed {len(qa_pairs)} QA pairs to {index_path}")

    return 0


if __name__ == "__main__":
    sync_data()